# utils/memory_system.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
记忆系统
"""
import json
import logging
import os
import re
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple

# 本地模块导入
try:
    from utils.llm_service import LLMService
except ImportError:
    LLMService = None  # type: ignore

from models import User, UserBuffer
from utils import llm_service

logger = logging.getLogger("AI-MindCare-System-Memory")


class MemorySystem:
    """记忆系统，管理用户信息、咨询记录、对话历史等"""
    def __init__(self, llm_service: Optional[LLMService] = None):
        self.llm_service = llm_service

    
    def _load_long_term_memory(self, user_id: uuid.UUID) -> Optional[Dict[str, Any]]:
        """加载用户的长期记忆"""
        user = User.query.get(user_id)
        if user and user.long_term_memory:
            return user.long_term_memory
        return None
    def _update_long_term_memory(self, user_id: uuid.UUID, long_term_memory: Dict[str, Any]):
        """更新用户的长期记忆"""
        user = User.query.get(user_id)
        if user:
            user.long_term_memory = long_term_memory
        else:
            new_user = User(id=user_id, long_term_memory=long_term_memory)
            db.session.add(new_user)
        try:
            db.session.commit()
            logger.info(f"用户 {user_id} 的长期记忆已更新到数据库")
        except Exception as e:
            db.session.rollback()
            logger.error(f"更新用户 {user_id} 长期记忆到数据库时出错: {e}")

    
    def _get_user_buffers(self) -> Dict[str, Dict[str, List[Any]]]:
        """从 UserBuffer 表获取所有用户的会话缓冲区"""
        buffers = {}
        records = UserBuffer.query.order_by(
            UserBuffer.conversation_id,
            UserBuffer.timestamp.desc()
        ).all()
        for record in records:
            user_id_str = str(record.user_id)
            if user_id_str not in buffers:
                buffers[user_id_str] = {}
            if record.conversation_id not in buffers[user_id_str]:
                buffers[user_id_str][record.conversation_id] = []
            buffers[user_id_str][record.conversation_id].append(record.message)
        return buffers


    def _get_user_buffer(self, conversation_id: str) -> List[str]:
        """获取用户某个会话的缓冲区"""
        records = UserBuffer.query.filter_by(
            conversation_id=conversation_id
        ).order_by(UserBuffer.timestamp.desc()).all()
        return [record.message for record in records]

    def add_to_user_buffer(self, user_id: uuid.UUID, conversation_id: str, message: str, role: str):
        """向用户的会话缓冲区添加新消息"""
        new_conversation_record = ConversationRecord(
            user_id=user_id,
            conversation_id=conversation_id,
            role=role,
            content=message,
            message_id=str(uuid.uuid4()),
            created_at=datetime.now(timezone.utc)
        )
        db.session.add(new_conversation_record)
        new_buffer = UserBuffer(user_id=user_id,
            conversation_id=conversation_id,
            role=role,
            content=message,
            message_id=str(uuid.uuid4()),
            created_at=datetime.now(timezone.utc))
        db.session.add(new_buffer)

        try:
            db.session.commit()
            # 清理 UserBuffer 表中旧记录，只保留每个会话最近的 20 条
            old_records = UserBuffer.query.filter_by(
                user_id=user_id,
                conversation_id=conversation_id
            ).order_by(UserBuffer.timestamp.asc()).limit(
                db.session.query(UserBuffer).filter_by(
                    user_id=user_id,
                    conversation_id=conversation_id
                ).count() - 20
            )
            for record in old_records:
                db.session.delete(record)
            db.session.commit()
            logger.info(f"已为用户 {user_id} 会话 {conversation_id} 添加新消息")
        except Exception as e:
            db.session.rollback()
            logger.error(f"为用户 {user_id} 会话 {conversation_id} 添加新消息时出错: {e}")





    def get_user_buffer(self, user_id: uuid.UUID, conversation_id: str) -> List[str]:
        """获取用户某个会话的缓冲区"""
        records = UserBuffer.query.filter_by(
            user_id=user_id,
            conversation_id=conversation_id
        ).order_by(UserBuffer.timestamp.desc()).limit(20).all()
        return [record.message for record in records]
######################################


    def _update_long_term_memory(self, user_id: int, long_term_memory: Dict[str, Any]):
        record = UserMemory.query.filter_by(user_id=user_id).first()
        if record:
            record.long_term_memory = long_term_memory
        else:
            new_record = UserMemory(user_id=user_id, long_term_memory=long_term_memory, conversation_id='', user_buffer=[])
            db.session.add(new_record)
        try:
            db.session.commit()
            logger.info(f"用户 {user_id} 的长期记忆已更新到数据库")
        except Exception as e:
            db.session.rollback()
            logger.error(f"更新用户 {user_id} 长期记忆到数据库时出错: {e}")

    def create_conversation_context(self, user_id: int, conversation_id: str, context: Dict[str, Any]):
        record = UserMemory.query.filter_by(user_id=user_id, conversation_id=conversation_id).first()
        if record:
            record.user_buffer = []
        else:
            new_record = UserMemory(user_id=user_id, conversation_id=conversation_id, user_buffer=[], long_term_memory={})
            db.session.add(new_record)
        try:
            db.session.commit()
            logger.info(f"为用户 {user_id} 会话 {conversation_id} 创建上下文成功")
        except Exception as e:
            db.session.rollback()
            logger.error(f"为用户 {user_id} 会话 {conversation_id} 创建上下文时出错: {e}")

    def get_conversation_context(self, user_id: int, conversation_id: str) -> Dict[str, Any]:
        record = UserMemory.query.filter_by(user_id=user_id, conversation_id=conversation_id).first()
        context = {}
        if record:
            context['user_buffer'] = record.user_buffer
            context.update(record.long_term_memory or {})
        return context

    def _get_user_conversation_history(self, user_id: int) -> Dict[str, Any]:
        records = UserMemory.query.filter_by(user_id=user_id).all()
        history = {}
        for record in records:
            if record.conversation_id:
                history[record.conversation_id] = {
                    'user_buffer': record.user_buffer,
                    'long_term_memory': record.long_term_memory
                }
        return history

    def _create_user_data(self, user_id: int):
        new_record = UserMemory(user_id=user_id, conversation_id='', user_buffer=[], long_term_memory={})
        try:
            db.session.add(new_record)
            db.session.commit()
            logger.info(f"为用户 {user_id} 创建数据成功")
        except Exception as e:
            db.session.rollback()
            logger.error(f"为用户 {user_id} 创建数据时出错: {e}")
########################################################################################
   
    def _now_utc(self) -> datetime:
        return datetime.now(timezone.utc)

    def _load_memory(self) -> Dict[str, Any]:
        """从文件中加载数据"""
        if self.persistence_path and os.path.exists(self.persistence_path):
            try:
                with open(self.persistence_path, 'r', encoding='utf-8') as f:
                    logger.info(f"从 {self.persistence_path} 加载记忆...")
                    loaded_data = json.load(f)
                    loaded_data.setdefault("users", {})
                    loaded_data.setdefault("metadata", {})
                    return loaded_data
            except (json.JSONDecodeError, IOError) as e:
                logger.error(
                    f"加载记忆文件 '{self.persistence_path}' 失败: {e}，将使用空记忆。")
        return {
            "users": {},
            "metadata": {},
        }

    def _save_memory(self):
        """将数据保存到文件"""
        if self.persistence_path:
            try:
                os.makedirs(os.path.dirname(self.persistence_path),
                            exist_ok=True)
                with open(self.persistence_path, 'w', encoding='utf-8') as f:
                    json.dump(self._memory,
                              f,
                              indent=2,
                              ensure_ascii=False,
                              default=str)
                logger.info(f"记忆已保存到 {self.persistence_path}")
            except IOError as e:
                logger.error(f"保存记忆文件 '{self.persistence_path}' 失败: {e}")
            except TypeError as e:
                logger.error(f"序列化记忆时失败: {e}", exc_info=True)

    def _get_user_conversation_history(self, user_id: str):
        """获取用户的完整对话历史(字典类型)"""
        user_data = self._memory.get("users", {}).get(user_id, {})
        return user_data.get("conversations", {})

    def _update_long_term_memory(self, user_id: str, new_memory: dict):
        """更新用户的长期记忆(字典类型)"""
        users = self._memory.get("users", {})
        user_record = users.setdefault(user_id, {})
        existing_memory = user_record.setdefault("long_term_memory", {})
        if not isinstance(existing_memory, dict):
            existing_memory = {}
        existing_memory.update(new_memory)
        user_record["long_term_memory"] = existing_memory
        users[user_id] = user_record
        self._memory["users"] = users
        self._save_memory()
        logger.info(f"用户 {user_id} 的长期记忆已更新。")

    def _load_long_term_memory(self, user_id: str):
        """加载用户的长期记忆(字典类型)"""
        return self._memory.get("users",
                                {}).get(user_id,
                                        {}).get("long_term_memory", {})

    def _get_user_buffers(self):
        """初始化用户对话缓冲区"""
        users = self._memory.get("users", {})
        user_buffers = {}
        for user_id, user_data in users.items():
            user_buffers[user_id] = {}
            conversation_history = user_data.get("conversations", {})
            for conversation_id, conversation in conversation_history.items():
                messages = conversation.get("messages", [])
                user_buffers[user_id][conversation_id] = messages[
                    -20:]  # 保留最近的20条对话信息
        return user_buffers

    def _create_user_data(self, user_id: str) -> Dict[str, Any]:
        """创建用户数据"""
        users = self._memory.setdefault("users", {})
        if user_id not in users:
            logger.info(f"用户 {user_id} 不存在，创建新的记忆记录。")
            users[user_id] = {
                "basic_info": {},
                "long_term_memory": {},
                "conversations": {},
                "last_updated": self._now_utc().isoformat()
            }
            self._memory["users"] = users
            self._save_memory()
        return users[user_id]

    def _get_user_data(
            self,
            user_id: str,
            create_if_not_exists: bool = True) -> Optional[Dict[str, Any]]:
        """获取用户数据（支持自动创建）"""
        users = self._memory.setdefault("users", {})
        if user_id not in users and create_if_not_exists:
            self._create_user_data(user_id)
        user_data = users.get(user_id, {})
        # 确保必要字段存在
        user_data.setdefault("basic_info", {})
        user_data.setdefault("long_term_memory", {})
        user_data.setdefault("conversations", {})
        user_data.setdefault("last_updated", self._now_utc().isoformat())
        self._memory["users"] = users
        self._save_memory()
        return user_data

    def add_or_update_user_info(self, user_id: str, basic_info: Dict[str,
                                                                     Any]):
        """更新用户基本信息"""
        user_data = self._get_user_data(user_id, create_if_not_exists=True)
        if user_data:
            user_data["basic_info"].update(basic_info)
            user_data["last_updated"] = self._now_utc().isoformat()
            self._save_memory()

    def get_user_info(self, user_id: str) -> Optional[Dict[str, Any]]:
        """获取用户基本信息"""
        user_data = self._get_user_data(user_id, create_if_not_exists=False)
        return user_data.get("basic_info") if user_data else None

    def get_all_users(self) -> Dict[str, Dict[str, Any]]:
        """获取所有用户信息"""
        return {
            uid: data.get("basic_info", {})
            for uid, data in self._memory.get("users", {}).items()
        }

    def get_all_user_ids(self) -> List[str]:
        """获取所有用户ID"""
        return list(self._memory.get("users", {}).keys())

    def create_conversation_context(self,
                                    user_id: str,
                                    conversation_id: str,
                                    context: Dict[str, Any] = None):
        """创建会话上下文（包含默认初始化字段）"""
        if not user_id or not conversation_id:
            logger.error("创建上下文失败：未提供 user_id 或 conversation_id")
            return

        # 初始化上下文结构
        initial_context = {
            "last_updated": self._now_utc().isoformat(),
            "topic": "",
            "messages": [],
            "status": "active"
        }
        if context:
            initial_context.update(context)

        conversations = self._memory.get("users",
                                         {}).get(user_id,
                                                 {}).get("conversations", {})
        if conversation_id not in conversations:
            conversations[conversation_id] = initial_context
            self._memory["users"][user_id]["conversations"] = conversations
            self._save_memory()
            logger.info(f"成功创建会话上下文：{conversation_id}（用户：{user_id}）")
        else:
            logger.warning(f"会话上下文已存在：{conversation_id}（用户：{user_id}）")

    def get_conversation_context(self, user_id: str,
                                 conversation_id: str) -> Dict[str, Any]:
        """获取会话上下文（不存在时返回空字典）"""
        if not user_id or not conversation_id:
            logger.error("获取上下文失败：未提供 user_id 或 conversation_id")
            return {}
        return self._memory.get("users",
                                {}).get(user_id,
                                        {}).get("conversations",
                                                {}).get(conversation_id, {})

    def update_conversation_context(self, user_id: str, conversation_id: str,
                                    updates: Dict[str, Any]):
        """更新会话上下文"""
        if not user_id or not conversation_id:
            logger.error("更新上下文失败：未提供 user_id 或 conversation_id")
            return

        conversations = self._memory.get("users",
                                         {}).get(user_id,
                                                 {}).get("conversations", {})
        if conversation_id in conversations:
            conversations[conversation_id].update(updates)
            conversations[conversation_id]["last_updated_utc"] = self._now_utc(
            ).isoformat()
            self._save_memory()
            logger.debug(f"更新会话上下文：{conversation_id}（用户：{user_id}）")
        else:
            logger.error(f"会话上下文不存在：{conversation_id}（用户：{user_id}）")

    def delete_conversation_context(self, user_id: str, conversation_id: str):
        """删除会话上下文"""
        if not user_id or not conversation_id:
            logger.error("删除上下文失败：未提供 user_id 或 conversation_id")
            return

        conversations = self._memory.get("users",
                                         {}).get(user_id,
                                                 {}).get("conversations", {})
        if conversation_id in conversations:
            del conversations[conversation_id]
            self._save_memory()
            logger.info(f"成功删除会话上下文：{conversation_id}（用户：{user_id}）")
        else:
            logger.warning(f"会话上下文不存在：{conversation_id}（用户：{user_id}）")

    def generate_topic_for_conversation(self, user_id: str,
                                        conversation_id: str) -> str:
        """为会话生成摘要（需配置 LLM 服务）"""
        logger.info(f"开始为会话 {conversation_id}（用户：{user_id}）生成摘要...")
        try:
            # 获取对话消息
            conversation = self.get_conversation_context(
                user_id, conversation_id)
            messages = conversation.get("messages", [])
            if not messages:
                return "会话中没有消息，无法生成摘要。"

            # 构建提示词
            prompt = "\n".join([
                f"{msg['role']} ({msg['timestamp']}): {msg['content']}"
                for msg in messages
            ])
            summary_prompt = f"""
            请对以下对话生成主题摘要：
            
            对话内容：
            {prompt}
            
            摘要要求：
            1. 简洁明了，仅包含关键信息。
            2. 确保摘要内容与对话内容相关。
            3. 避免重复信息。
            4. 保持语言流畅。
            5. 确保摘要能够准确反映对话主题和主要内容。
            6. 摘要长度不超过20字。

            输出示例：
            * 产品研发技术路线规划
            * 解释CGCS2000大地坐标系
            * Python爬取网站数据
            """

            if not self.llm_service:
                logger.error("LLM 服务未初始化，无法生成摘要")
                return "服务暂不可用，请稍后再试。"

            # 调用大模型生成会话的主题
            topic = self.llm_service.generate_response(
                prompt=summary_prompt,
                system_message="你是专业的心理咨询摘要生成器，需用简洁语言提炼关键信息",
                max_tokens=500,
                temperature=0.5)
            logger.info(f"成功生成会话主题摘要：{conversation_id}（用户：{user_id}）")
            conversation["topic"] = topic.strip()
            self.update_conversation_context(user_id, conversation_id,
                                             {"topic": topic.strip()})
            self._save_memory()
            return topic

        except Exception as e:
            logger.error(f"生成摘要失败：{str(e)}", exc_info=True)
            return "生成摘要时发生错误，请检查日志"
